#include "llvm/Module.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Constants.h"
#include "llvm/BasicBlock.h"
#include "llvm/Pass.h"
#include "llvm/Support/CFG.h"
#include "llvm/Support/CallSite.h"
#include "llvm/IntrinsicInst.h"
#include "llvm/LinkAllVMCore.h"
#include "llvm/CodeGen/DwarfWriter.h"
#include "llvm/TypeSymbolTable.h"
#include "llvm/Bitcode/ReaderWriter.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/LLVMContext.h"

#include "trace-sync-events.h"
#include "../../../include/tern/TernDefs.h"
#include "../../../include/tern/TernLandmark.h"
#include <ext/hash_map>

#define UI64_TYPE IntegerType::getInt64Ty(getGlobalContext())
#define UI32_TYPE IntegerType::getInt32Ty(getGlobalContext())

using namespace std;
using namespace llvm;
/*
   struct eqstr {
   bool operator()(const char* s1, const char* s2) const {
   return strcmp(s1, s2) == 0;
   }
   };
   */
typedef __gnu_cxx::hash_map<const char*, llvm::Function*, 
		__gnu_cxx::hash<const char*>, eqstr> func_map;

namespace defens {
	struct trace_sync_events_stub: public FunctionPass {
		uint64_t hookID;
		static char ID; // Pass identification

		trace_sync_events_stub(): FunctionPass(&ID) {}            
		virtual bool doInitialization(Module &M);
		virtual bool doFinalization(Module &M);
		//virtual bool doFinalization(Function &F);
		//virtual bool runOnBasicBlock(BasicBlock &BB);
		virtual bool runOnFunction(Function &F);
		virtual bool updateInstIds(BasicBlock &BB);
		/*virtual bool doInitialization(Function &F) {
		  return FunctionPass::doInitialization(F);
		  }*/
		Function *isSyncEventFunction(Instruction *Inst, uint64_t &funcIndex);
		int getFixOpInfo(Instruction *Inst, struct FixOpInfo *Info);
		string getAnnotation(Instruction *Inst);
		bool DFSForward(Instruction *Inst, string &FixID, int FixType, bool IsFirst, int Depth);
		bool DFSBackward(Instruction *Inst, string &FixID, int FixType, bool IsFirst, int Depth);
		void pushSyncVars(Function *F, CallSite &CS, std::vector<Value*> &args);
		void buildBeforeHook(Function *F, CallSite &CS, Instruction *Inst, uint64_t funcIndex);
		void buildAfterHook(Function *F, CallSite &CS, Instruction *Inst, uint64_t funcIndex);
		Value* getArgIntTypeValue(Instruction *Inst, Value *V);
		void pushCallSiteArgs(Function *F, CallSite &CS, Instruction *Inst, std::vector<Value*> &Args);
		void init_funcs(Module &M);
		func_map funcs;
		// Jingyue: the declaration of <tern_capture_input>
		// Assigned in doInitialization(Module &M)
		Function *ternCaptureInput;
	};

	void trace_sync_events_stub::init_funcs(Module &M) {

		// initialize set of functions to load
		for(unsigned i=0; i<(sizeof landmarks)/(sizeof landmarks[0]); ++i) {
			//if (isIgnoredLandmarkAtInsertion(i))
			//    continue;
			funcs[landmarks[i]] = NULL;
		}

		// get template functions from template.bc
		string template_path(getenv("DEFENS_ROOT"));
		template_path += "/rcs-modules/klee/llvm-instrument//trace-sync-events/stub/template.bc";
		Module *dummy = ParseBitcodeFile(MemoryBuffer::getFile(template_path), getGlobalContext());
		for(Module::iterator fi=dummy->begin(); fi!=dummy->end(); ++fi) {
			if(fi->getName() != "template")
				continue;
			// std::cerr << "function name: " << fi->getNameStr() << endl;
			for(Function::iterator bi=fi->begin(); bi!=fi->end(); ++bi) {
				for(BasicBlock::iterator ii=bi->begin(); ii!=bi->end(); ++ii){
					if(!isa<CallInst>(ii))
						continue;
					CallInst *c = dyn_cast<CallInst>(ii);
					Function *f = c->getCalledFunction();
					//std::cerr << *ii << endl;
					if(funcs.find(f->getNameStr().c_str()) != funcs.end()) {
						//std::cerr << "loading " << f->getNameStr() << endl;
						funcs[f->getNameStr().c_str()] = f;
					}
				}
			}
		}

		// insert functions to target module
		for(func_map::iterator it=funcs.begin(); it!=funcs.end(); ++it) {
			// Heming: only insert those functions used in template.bc.
			if (it->second != NULL) {
				const char* name = it->first;
				//std::cerr << "inserting " << name << endl;
				assert(funcs[name] && "function is not loaded!");
				// const FunctionType *type = funcs[name]->getFunctionType();
				// std::cerr << *type << endl;
				Constant *c = M.getOrInsertFunction(name, funcs[name]->getFunctionType());
				Function *f = dyn_cast<Function>(c);
				//std::cerr << "Inserted " << *f << endl;
				funcs[name] = f;
			}
		}
	}

	bool trace_sync_events_stub::doInitialization(Module &M) {
		hookID = 0;
		init_funcs(M);
		// insert tern_init at the entry of main
		for(Module::iterator fi=M.begin(); fi!=M.end(); ++fi) {
			if (fi->getName() != "main")
				continue;
			Instruction *pos = fi->getEntryBlock().getFirstNonPHI();
			vector<Value*> params;
			for (Function::arg_iterator ai = fi->arg_begin(); ai != fi->arg_end(); ai++) {
				if (params.size() >= 2)
					break;
				params.push_back(ai);
			}
			CallInst::Create(funcs["tern_init"], params.begin(), params.end(), "", pos);
		}
		// create the declaration of <tern_capture_input>
		// Set up types
		const Type *charType = IntegerType::get(getGlobalContext(), 8);
		const Type *strType = PointerType::getUnqual(charType);
		const Type *voidType = Type::getVoidTy(getGlobalContext());
		vector<const Type *> paramTypes;
		paramTypes.push_back(strType);
		paramTypes.push_back(UI32_TYPE);
#if __WORDSIZE == 64
		paramTypes.push_back(UI64_TYPE);
#else
		paramTypes.push_back(UI32_TYPE);
#endif
		const FunctionType *funcType = FunctionType::get(voidType, paramTypes, false); 
		// Set up the function <tern_capture_input>
		ternCaptureInput = dyn_cast<Function>(M.getOrInsertFunction(TERN_CAPTURE_INPUT, funcType));
		return true;
	}

	bool trace_sync_events_stub::doFinalization(Module &M) {
		/* Heming: this is important, since we have inserted hook functions, the instruction ids may have been changed,
		   so we need to reparse the module to update instruction ids passed to hook functions. 
		   Here the assumption is if we only change the argument of hook function call sites, the inst ids will not change.
		   Need to ask Jingyue to make sure about this. */
		for (Module::iterator f = M.begin(), fe = M.end(); f != fe; ++f)
			for (Function::iterator b = f->begin(), be = f->end(); b != be; ++b)
				updateInstIds(*b);

		return true;
	}

	/*bool trace_sync_events_stub::doFinalization(Function &F) {
	  return true;
	  }*/

	Function *trace_sync_events_stub::isSyncEventFunction(Instruction *Inst, uint64_t &funcIndex) {
		Function *F = NULL;
		int index = -1;

		if (Inst->getOpcode() == Instruction::Call)
			F = dyn_cast<CallInst>(Inst)->getCalledFunction();
		else if (Inst->getOpcode() == Instruction::Invoke)
			F = dyn_cast<InvokeInst>(Inst)->getCalledFunction();
		if (F == NULL)
			return NULL;
		const char *funcName = F->getNameStr().c_str();
		if (F && (index = getLandmarkIndex(funcName)) >= 0) {
			funcIndex = (uint64_t)index;
			return F;
		}
		return NULL;
	}

	void trace_sync_events_stub::pushSyncVars(Function *F, CallSite &CS, std::vector<Value*> &args) {
		StringRef condSignalName("pthread_cond_signal");
		Function *condSignalFunc = F->getParent()->getFunction(condSignalName);
		const Type *condType = condSignalFunc->arg_begin()->getType();
		assert(condType);

		StringRef MutexName("pthread_mutex_lock");
		Function *mutexLock = F->getParent()->getFunction(MutexName);
		const Type *mutexType = mutexLock->arg_begin()->getType();
		assert(mutexType);

		if (F->getNameStr() == "pthread_cond_wait" || F->getNameStr() == "pthread_cond_timedwait") {
			fprintf(stderr, "instrument pthread wait functions\n");
			args.push_back(CS.getArgument(0));
			args.push_back(CS.getArgument(1));
		} else if (F->getNameStr() == "pthread_cond_signal" || F->getNameStr() == "pthread_cond_broadcast") {
			fprintf(stderr, "instrument pthread signal functions\n");
			args.push_back(CS.getArgument(0));
			args.push_back(Constant::getNullValue(mutexType));
		} else if (F->getNameStr() == "pthread_mutex_lock" || F->getNameStr() == "pthread_mutex_unlock" 
				|| F->getNameStr() == "pthread_mutex_trylock") {
			fprintf(stderr, "instrument pthread mutex functions\n");
			args.push_back(Constant::getNullValue(condType));
			args.push_back(CS.getArgument(0));
		} else {
			args.push_back(Constant::getNullValue(condType));
			args.push_back(Constant::getNullValue(mutexType));
		}
		fprintf(stderr, "instrument pthread functions finish\n");
	}

	Value* trace_sync_events_stub::getArgIntTypeValue(Instruction *Inst, Value *V) {
		if (isa<PointerType>(V->getType())) {
			//fprintf(stderr, "getArgIntTypeValue PointerType\n");
			return CastInst::CreatePointerCast(V, UI64_TYPE, V->getName(), Inst);
		}
		else if (isa<IntegerType>(V->getType())) {
			//fprintf(stderr, "getArgIntTypeValue IntegerType\n");
			return CastInst::CreateIntegerCast(V, UI64_TYPE, false, V->getName(), Inst);
		}
		else {
			fprintf(stderr, "getArgIntTypeValue Unknown Type\n");
			assert(false);
		}
		return 0;
	}

	// Push up to three arguments in the call site so far.
	void trace_sync_events_stub::pushCallSiteArgs(Function *F, CallSite &CS, Instruction *Inst, std::vector<Value*> &Args) {
		// default value, if no argument.
		APInt Zero(64, 0);

		if (CS.arg_size() >= 1) {
			CallSite::arg_iterator AI = CS.arg_begin();
			Value *V = *AI;
			Args.push_back(getArgIntTypeValue(Inst, V));

			if (CS.arg_size() >= 2) {
				AI++;
				V = *AI;
				Args.push_back(getArgIntTypeValue(Inst, V));

				if (CS.arg_size() >= 3) {
					AI++;
					V = *AI;
					Args.push_back(getArgIntTypeValue(Inst, V));
				} else {
					Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
				}
			} else {
				Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
				Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
			}
		} else {
			Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
			Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
			Args.push_back(Constant::getIntegerValue(UI64_TYPE, Zero));
		}		
	}

	// void tern_before_landmark(unsigned landmark_type, unsigned hook_id, uint64_t arg1, uint64_t arg2, uint64_t arg3);
	void trace_sync_events_stub::buildBeforeHook(Function *F, CallSite &CS, Instruction *Inst, uint64_t funcIndex) {
		std::vector<Value*> beforeHookArgs;
		StringRef beforeHookName(TERN_BEFORE_LANDMARK);
		Function *beforeHookFunc = F->getParent()->getFunction(beforeHookName);
		assert(beforeHookFunc);

		// unsigned landmark_type
		APInt AI(32, funcIndex);
		beforeHookArgs.push_back(Constant::getIntegerValue(UI32_TYPE, AI));

		// unsigned hook_id
		APInt HookID(32, hookID);
		beforeHookArgs.push_back(Constant::getIntegerValue(UI32_TYPE, HookID));

		// uint64_t arg1, uint64_t arg2, uint64_t arg3
		pushCallSiteArgs(beforeHookFunc, CS, Inst, beforeHookArgs);

		// insert
		//fprintf(stderr, "before hook arg num %u\n", beforeHookArgs.size());
		CallInst::Create(beforeHookFunc, beforeHookArgs.begin(), beforeHookArgs.end(), "", Inst);
	}

	// void tern_after_landmark(pthread_t *tid, unsigned landmark_type, unsigned hook_id, uint64_t ret_val, uint64_t arg1, uint64_t arg2, uint64_t arg3);
	void trace_sync_events_stub::buildAfterHook(Function *F, CallSite &CS, Instruction *Inst, uint64_t funcIndex) {
		std::vector<Value*> afterHookArgs;
		StringRef afterHookName(TERN_AFTER_LANDMARK);
		Function *afterHookFunc = F->getParent()->getFunction(afterHookName);
		assert(afterHookFunc);

		// pthread_t *tid
		if (F->getNameStr() == "pthread_create") {
			afterHookArgs.push_back(CS.getArgument(0));
		} else {
#if __WORDSIZE == 64
			afterHookArgs.push_back(Constant::getNullValue(PointerType::getUnqual(UI64_TYPE)));
#else
			afterHookArgs.push_back(Constant::getNullValue(PointerType::getUnqual(UI32_TYPE)));
#endif
		}

		// unsigned landmark_type
		APInt AI(32, funcIndex);
		afterHookArgs.push_back(Constant::getIntegerValue(UI32_TYPE, AI));

		// unsigned hook_id
		APInt HookID(32, hookID);
		afterHookArgs.push_back(Constant::getIntegerValue(UI32_TYPE, HookID));

		// uint64_t ret_val
		if (F->getReturnType() != Type::getVoidTy(getGlobalContext())) {
			afterHookArgs.push_back(getArgIntTypeValue(Inst, CS.getInstruction()));
		} else {
			afterHookArgs.push_back(Constant::getIntegerValue(UI64_TYPE, APInt(64, 0)));
		}

		// uint64_t arg1, uint64_t arg2, uint64_t arg3
		pushCallSiteArgs(afterHookFunc, CS, Inst, afterHookArgs);

		//fprintf(stderr, "after hook arg num %u\n", afterHookArgs.size());
		CallInst::Create(afterHookFunc, afterHookArgs.begin(), afterHookArgs.end(), "", Inst);

	}

	bool trace_sync_events_stub::runOnFunction(Function &F) {
		Function *FP = NULL;
		uint64_t funcIndex;
		CallSite CS;

		std::ofstream hookIdFunctionLog;
		char buf[BASIC_SIZE];
		memset(buf, 0, BASIC_SIZE);
		sprintf(buf, "./hookfunc-%ld.log", time(NULL));
		hookIdFunctionLog.open(buf, ios_base::app);

		for (Function::iterator bb = F.begin(), be = F.end(); bb != be; ++bb) {
			BasicBlock::iterator IB = bb->begin();
			BasicBlock::iterator IE = bb->end();

			// For invoke instruction at the end of the previous basic block, inser after hook to the beginning of the current basic block.
			// This happens for pthread_cond_wait() in PBZip2.
			if (FP) {
				buildAfterHook(FP, CS, IB, funcIndex);
				FP = NULL;
			}

			while (IB != IE) {
				BasicBlock::iterator NIB = IB;
				NIB++;
				if ((FP = isSyncEventFunction(IB, funcIndex)) != NULL && !isIgnoredLandmarkAtInsertion(funcIndex)) {			
					hookID++;
					string funcName = landmarks[funcIndex];
					hookIdFunctionLog << hookID << ":     				" << funcName
						<< ":				Function: " <<  F.getNameStr() 
						<< ",				BB: " << bb->getNameStr() << "\n"; 

					//Insert two instructions, one is calling the subroutine directly, the other is exit().
					//fprintf(stderr, "Function %s, index %llu, hookID %llu.\n", F->getNameStr().c_str(), funcIndex, hookID);
					//std::cerr << "Inst: " << *IB << "\n";
					CS = CallSite::get(IB);
					BasicBlock::iterator BI = IB;
					buildBeforeHook(FP, CS, BI, funcIndex);
					BI++;
					if (BI == bb->end()) {
						// Do not insert after hook now, but at the beginning of the next basic block.
						break;
					}
					buildAfterHook(FP, CS, BI, funcIndex);
					FP = NULL;			// This reset is necessary.
				}
				IB = NIB;
			}
		}

		hookIdFunctionLog.flush();
		hookIdFunctionLog.close();

		// Jingyue: In order to capture the inputs to httpd, 
		// add <tern_capture_input> to the end of <ap_rgetline_core>.
		
		if (F.getNameStr() == "ap_rgetline_core") {
			// Set up the input, which is the formal parameter of the function
			assert(F.arg_size() >= 3);
			Value *arg0 = F.arg_begin();
			Value *arg2 = ++(++F.arg_begin());
			// Add the function <tern_capture_input> before each return instruction.
			assert(ternCaptureInput && "Cannot find or insert <tern_capture_input>");
			for (Function::iterator bi = F.begin(); bi != F.end(); ++bi) {
				ReturnInst *ri = dyn_cast<ReturnInst>(bi->getTerminator());
				if (ri) {
					Value *input = new LoadInst(arg0, "input", ri);
					Value *retValue = ri->getReturnValue();
					assert(retValue && "The function does not return a value");
					Value *read = new LoadInst(arg2, "read", ri);
					vector<Value *> params;
					params.push_back(input);
					params.push_back(retValue);
					params.push_back(read);
					CallInst::Create(ternCaptureInput, params.begin(), params.end(), "", ri);
				}
			}
		}

		return true;
	}

	bool trace_sync_events_stub::updateInstIds(BasicBlock &BB) {
		// For each landmark call site, first get the inst id, and then trace backward and forward within the basicblock,
		// update the inst id passed to before and after landmark hook function.

		return true;
	}

	char trace_sync_events_stub::ID = 0;
	static RegisterPass<trace_sync_events_stub> X("trace_sync_events_stub", "trace sync events stub");
}

